from __future__ import division
import numpy as np
import sys
import time

#Initial idea from David Warde-Farley on the SciPy Cookbook
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle
from matplotlib.ticker import NullLocator
import warnings
import time
from termcolor import colored
warnings.simplefilter("ignore", np.ComplexWarning)

class Computation_time():
    """
    Not done for recursive algorithm
    """
    def __init__(self,enabled=True):
        self.enabled=enabled
        self.Max={}
        self.Average={}
        self.Occurences={}
        self.Total={}
        self.Start_time={}
        self.Level={}
        self.Last_duration={}
    
    def TIC(self,function_name="test",level="L"):
        if self.enabled:
            self.Start_time[function_name]=time.time()
            self.Level[function_name]=level
    
    def TOC(self,function_name="test"):
        if self.enabled:
            toc=time.time()
            if function_name in self.Total:
                duration=toc-self.Start_time[function_name]
                self.Total[function_name]+=duration
                n=self.Occurences[function_name]
                self.Average[function_name]=(self.Average[function_name]*n+duration)/(n+1)
                self.Occurences[function_name]+=1
                if self.Max[function_name]<duration:
                    self.Max[function_name]=duration
                self.Last_duration[function_name]=duration
            else:
                duration=toc-self.Start_time[function_name]
                self.Total[function_name]=duration
                self.Average[function_name]=duration
                self.Occurences[function_name]=1
                self.Max[function_name]=duration
                self.Last_duration[function_name]=duration
    
    def Calculate_ratios(self,fi1,fi2,ff,level):
            self.Level[ff]="R"
            if ff in self.Average:
                ratio=self.Last_duration[fi1]/self.Last_duration[fi2]
                n=self.Occurences[ff]
                self.Average[ff]=(self.Average[ff]*n+ratio)/(n+1)
                self.Occurences[ff]+=1
                if self.Max[ff]<ratio:
                    self.Max[ff]=ratio
            else:
                ratio=self.Last_duration[fi1]/self.Last_duration[fi2]
                self.Average[ff]=ratio
                self.Occurences[ff]=1
                self.Max[ff]=ratio
    
    def Display_ratios(self,level):   
        for key in self.Average:
            if self.Level[key] == "R":
                print colored(key,'green'),": ",colored(self.Occurences[key],'red'), " times, average= ",colored(self.Average[key],'blue'), " s, max= ",colored(self.Max[key],'blue')
        print "Done!"
        
    def Display(self,level="H"):
        if self.enabled:
            print ""
            if level=="L":
                lev="LMH"
            elif level=="M":
                lev="MH"
            else:
                lev="H"
            for key in self.Total:
                if self.Level[key] in lev:
                    print colored(key,'green'),": ",colored(self.Occurences[key],'red'), " times, total= ",colored(self.Total[key],'red'), " s, average= ",colored(self.Average[key],'blue'), " s, max= ",colored(self.Max[key],'blue')
            print "Done!"
            print ""
        else:
            print "Computation time recorder not enabled."

def Print_progress(current,total,start=0,pre_sentence='Progress: ',post_sentence=""):
    sentence='\r' + pre_sentence + '{:.1%}'.format((current-start)/(total-start))+post_sentence
    sys.stdout.write(sentence)
    sys.stdout.flush()
    
def Arange(ti,tf,step):
    return np.arange(ti,tf+step,step)

def hinton(w, maxWeight=None, ax=None):
    """
    Draws a Hinton diagram for visualizing a weight matrix.
    """
    W=w.T
    if not ax:
        fig = plt.figure()
        ax = fig.add_subplot(1, 1, 1)

    if not maxWeight:
        maxWeight = 2**np.ceil(np.log(np.abs(W).max())/np.log(2))

    ax.patch.set_facecolor('gray')
    ax.set_aspect('equal', 'box')
    ax.xaxis.set_major_locator(NullLocator())
    ax.yaxis.set_major_locator(NullLocator())

    for (x,y),w in np.ndenumerate(W):
        if w > 0: color = 'blue'
        else:     color = 'red'
        size = np.sqrt(np.abs(w))
        rect = Rectangle([x - size / 2, y - size / 2], size, size,
            facecolor=color, edgecolor=color)
        ax.add_patch(rect)
    ax.autoscale_view()

    # Reverse the axis limits
    #ax.set_xlim(*ax.get_xlim()[::-1])
    ax.set_ylim(*ax.get_ylim()[::-1])
    plt.title('Hinton')
    plt.show()

def null(A, eps=1e-15):
    u, s, vh = np.linalg.svd(A)
    null_space = np.compress(s <= eps, vh, axis=0)
    return null_space.T

def AVmult(a, v, eps=1e-15):
    res = np.round(np.matrix(a) * (v))
    if np.abs(res) <= eps and np.abs(res) >= -eps:
        return 0
    else:
        return res

def Round(a, eps=1e-16):
    res = np.abs(a)
    if res <= eps and res >= -eps:
        return 0
    else:
        return a


def FindIndex(v, l, eps=1e-12):
    res = 1
    for i in range(np.size(l)):
        if np.abs(l[i] - v) <= eps:
            res = i
    return res

def prod2(A, B):
    """
    Input:
        -A: first tensor (A=[A0,A1,A2])
        -B: second tensor (B=[B0,B1,B2])
    Outout:
        -res: dot(A_i,B_i)
    """
    res = []
    for i in [0, 1, 2]:
        res.append(np.kron(A[i], B[i]))
    return res

def scalprod2(A, B):
    """
    Input:
        -A: first tensor (A=[A0,A1,A2])
        -B: second tensor (B=[B0,B1,B2])
    Outout:
        -res: sum_i dot(A_i,B_i)
    """
    res = 0
    inter = prod2(A, B)
    for i in [0, 1, 2]:
        res += inter[i]
    return res
